# from visualizer import get_local
import torch
import torchinfo
import torch.nn as nn
from spikingjelly.clock_driven.neuron import (
    MultiStepParametricLIFNode,
    MultiStepLIFNode,
)
from spikingjelly.clock_driven import layer
from timm.models.layers import to_2tuple, trunc_normal_, DropPath
from timm.models.registry import register_model
from timm.models.vision_transformer import _cfg
from einops.layers.torch import Rearrange
import torch.nn.functional as F
from functools import partial


class BNAndPadLayer(nn.Module):# 自定义的批量归一化+填充层
    def __init__(
        self,
        pad_pixels,
        num_features,
        eps=1e-5,
        momentum=0.1,
        affine=True,
        track_running_stats=True,
    ):
        super(BNAndPadLayer, self).__init__()
        self.bn = nn.BatchNorm2d(
            num_features, eps, momentum, affine, track_running_stats
        )
        self.pad_pixels = pad_pixels
    
    # 本质还是批量归一化，变化是当affine=true时进行填充像素
    def forward(self, input):
        output = self.bn(input)
        if self.pad_pixels > 0:
            if self.bn.affine:
                pad_values = (
                    self.bn.bias.detach()
                    - self.bn.running_mean
                    * self.bn.weight.detach()
                    / torch.sqrt(self.bn.running_var + self.bn.eps)
                )
            else:
                pad_values = -self.bn.running_mean / torch.sqrt(
                    self.bn.running_var + self.bn.eps
                )
            output = F.pad(output, [self.pad_pixels] * 4)# 四个维度都填充pad_pixels个像素
            pad_values = pad_values.view(1, -1, 1, 1)# 调整形状
            output[:, :, 0 : self.pad_pixels, :] = pad_values
            output[:, :, -self.pad_pixels :, :] = pad_values
            output[:, :, :, 0 : self.pad_pixels] = pad_values
            output[:, :, :, -self.pad_pixels :] = pad_values
        return output
    # 这些property都是属性代理，反映了batchnorm2d的属性，在forward内调用并用于计算具体的像素数
    @property
    def weight(self):
        return self.bn.weight

    @property
    def bias(self):
        return self.bn.bias

    @property
    def running_mean(self):
        return self.bn.running_mean

    @property
    def running_var(self):
        return self.bn.running_var

    @property
    def eps(self):
        return self.bn.eps


class RepConv(nn.Module):
    def __init__(
        self,
        in_channel,
        out_channel,
        bias=False,
    ):
        super().__init__()
        # hidden_channel = in_channel
        conv1x1 = nn.Conv2d(in_channel, in_channel, 1, 1, 0, bias=False, groups=1)# 逐点卷积
        bn = BNAndPadLayer(pad_pixels=1, num_features=in_channel)
        conv3x3 = nn.Sequential(
            nn.Conv2d(in_channel, in_channel, 3, 1, 0, groups=in_channel, bias=False),
            nn.Conv2d(in_channel, out_channel, 1, 1, 0, groups=1, bias=False),
            nn.BatchNorm2d(out_channel),
        )

        self.body = nn.Sequential(conv1x1, bn, conv3x3)

    def forward(self, x):
        return self.body(x)


class SepConv(nn.Module):# 倒置分离卷积，逐点+深度+逐点
    r"""
    Inverted separable convolution from MobileNetV2: https://arxiv.org/abs/1801.04381.
    """

    def __init__(
        self,
        dim,# 输入/输出通道数
        expansion_ratio=2,# 中间层通道数相对于输入通道数的倍数
        act2_layer=nn.Identity,
        bias=False,
        kernel_size=7,
        padding=3,
    ):
        super().__init__()
        med_channels = int(expansion_ratio * dim)# 中间层通道数被放大2倍
        self.lif1 = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        self.pwconv1 = nn.Conv2d(dim, med_channels, kernel_size=1, stride=1, bias=bias)# 逐点卷积核心特点就是1*1的卷积核
        self.bn1 = nn.BatchNorm2d(med_channels)
        self.lif2 = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        self.dwconv = nn.Conv2d(# 当group=med_channels时，对每个通道单独卷积
            med_channels,
            med_channels,
            kernel_size=kernel_size,
            padding=padding,
            groups=med_channels,
            bias=bias,
        )  # depthwise conv
        self.pwconv2 = nn.Conv2d(med_channels, dim, kernel_size=1, stride=1, bias=bias)
        self.bn2 = nn.BatchNorm2d(dim)

    def forward(self, x):
        T, B, C, H, W = x.shape
        x = self.lif1(x)
        x = self.bn1(self.pwconv1(x.flatten(0, 1))).reshape(T, B, -1, H, W)
        x = self.lif2(x)
        x = self.dwconv(x.flatten(0, 1))
        x = self.bn2(self.pwconv2(x)).reshape(T, B, -1, H, W)
        return x


class MS_ConvBlock(nn.Module):# cnn的meta模块，由Sepconv和mlp组成
    def __init__(
        self,
        dim,
        mlp_ratio=4.0,
    ):
        super().__init__()

        self.Conv = SepConv(dim=dim)
        # self.Conv = MHMC(dim=dim)

        self.lif1 = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        self.conv1 = nn.Conv2d(
            dim, dim * mlp_ratio, kernel_size=3, padding=1, groups=1, bias=False
        )
        # self.conv1 = RepConv(dim, dim*mlp_ratio)
        self.bn1 = nn.BatchNorm2d(dim * mlp_ratio)  # 这里可以进行改进
        self.lif2 = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        self.conv2 = nn.Conv2d(
            dim * mlp_ratio, dim, kernel_size=3, padding=1, groups=1, bias=False
        )
        # self.conv2 = RepConv(dim*mlp_ratio, dim)
        self.bn2 = nn.BatchNorm2d(dim)  # 这里可以进行改进

    def forward(self, x):
        T, B, C, H, W = x.shape

        x = self.Conv(x) + x# 经过了所谓的token-mixer，也就是Sepconv
        x_feat = x
        x = self.bn1(self.conv1(self.lif1(x).flatten(0, 1))).reshape(T, B, 4 * C, H, W)
        x = self.bn2(self.conv2(self.lif2(x).flatten(0, 1))).reshape(T, B, C, H, W)
        x = x_feat + x

        return x


class MS_MLP(nn.Module):# Transformer-based SNN Block 的mlp部分
    def __init__(
        self, in_features, hidden_features=None, out_features=None, drop=0.0, layer=0
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        # self.fc1 = linear_unit(in_features, hidden_features)
        self.fc1_conv = nn.Conv1d(in_features, hidden_features, kernel_size=1, stride=1)
        self.fc1_bn = nn.BatchNorm1d(hidden_features)
        self.fc1_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")

        # self.fc2 = linear_unit(hidden_features, out_features)
        self.fc2_conv = nn.Conv1d(
            hidden_features, out_features, kernel_size=1, stride=1
        )  # 本质是nn.Linear(in_features, hidden_features)
        self.fc2_bn = nn.BatchNorm1d(out_features)
        self.fc2_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        # self.drop = nn.Dropout(0.1)

        self.c_hidden = hidden_features
        self.c_output = out_features

    def forward(self, x):
        T, B, C, H, W = x.shape
        N = H * W
        x = x.flatten(3)# (T, B, C, N)

        x = self.fc1_lif(x)
        x = self.fc1_conv(x.flatten(0, 1))
        x = self.fc1_bn(x).reshape(T, B, self.c_hidden, N).contiguous()

        x = self.fc2_lif(x)
        x = self.fc2_conv(x.flatten(0, 1))
        x = self.fc2_bn(x).reshape(T, B, C, H, W).contiguous()

        return x

# new
class DeltaEncoder(nn.Module):
    def __init__(self, dim):
        """
        简化的DeltaEncoder，输入输出都保持(T,B,C,H,W)的形状
        
        参数:
            dim: 输入特征的通道数
        """
        super().__init__()
        self.norm = nn.BatchNorm2d(dim)
        self.delta_proj = nn.Conv2d(dim, dim, kernel_size=1, padding=1)
        self.lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        
    def forward(self, x):
        """
        计算时间差分并输出与输入相同形状的特征
        
        参数:
            x: 形状为(T,B,C,H,W)的输入张量
            
        返回:
            形状为(T,B,C,H,W)的输出张量
        """
        # 输入 x: [T,B,C,H,W]
        T, B, C, H, W = x.shape
        
        # 计算时间差分
        delta = torch.zeros_like(x)
        delta[1:] = x[1:] - x[:-1]  # 第一帧的差分为0
        
        # 处理时间差分特征
        delta_flat = delta.flatten(0, 1)  # [T*B,C,H,W]
        delta_proj = self.delta_proj(delta_flat)
        delta_norm = self.norm(delta_proj)  # 标准化差分
        # delta_proj = self.delta_proj(delta_norm)  # 通过卷积层投影
        delta_reshape = delta_norm.reshape(T, B, C, H, W)  # 恢复原始形状
        
        # 通过LIF神经元处理
        delta_spk = self.lif(delta_reshape)  # 产生脉冲
        
        return delta_spk

# new change
class GRU_Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0, sr_ratio=1):
        super().__init__()
        assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
        self.dim = dim
        self.encoder = DeltaEncoder(dim=dim)
        # 简化版DeltaEncoder
        # self.delta_proj = nn.Conv2d(dim, dim, kernel_size=1)  # 1x1替代3x3
        # self.delta_norm = nn.BatchNorm2d(dim)
        
        # GRU相关组件
        # 更新门
        self.update_gate_conv = nn.Conv2d(dim*2, dim, kernel_size=3, padding=1)
        self.update_gate_bn = nn.BatchNorm2d(dim)
        # 重置门
        self.reset_gate_conv = nn.Conv2d(dim*2, dim, kernel_size=3, padding=1) 
        self.reset_gate_bn = nn.BatchNorm2d(dim)
        # 候选隐状态
        self.candidate_conv = nn.Conv2d(dim*2, dim, kernel_size=3, padding=1)
        self.candidate_bn = nn.BatchNorm2d(dim)
        
        # 输出投影 - 维持不变因为它已经是1x1卷积
        self.gru_proj = nn.Conv2d(dim, dim, kernel_size=1)
        self.gru_norm = nn.BatchNorm2d(dim)
        
        # 脉冲神经元
        self.lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        
    def forward(self, x):
        T, B, C, H, W = x.shape
        
        delta_features = self.encoder(x)
        # 简化的时间差分计算
        # delta = torch.zeros_like(x)
        # delta[1:] = x[1:] - x[:-1]  # 第一帧的差分为0
        # delta_flat = delta.flatten(0, 1)
        # delta_features = self.delta_norm(self.delta_proj(delta_flat)).reshape(T, B, C, H, W)
        # delta_features = self.lif(delta_features)
        
        # 初始化隐状态
        h_state = torch.zeros_like(x[0:1])  # [1,B,C,H,W]
        outputs = []
        
        for t in range(T):
            # 当前输入
            x_t = x[t:t+1]  # [1,B,C,H,W]
            delta_t = delta_features[t:t+1]  # [1,B,C,H,W]
            
            # 准备GRU输入 - 连接当前帧和隐状态
            gru_input = torch.cat([x_t, h_state], dim=2)  # [1,B,2C,H,W]
            gru_input_flat = gru_input.flatten(0, 1)  # [B,2C,H,W]
            
            # 计算更新门
            z_t = self.update_gate_conv(gru_input_flat)
            z_t = self.update_gate_bn(z_t)
            z_t = torch.sigmoid(z_t).reshape(1, B, C, H, W)
            
            # 计算重置门
            r_t = self.reset_gate_conv(gru_input_flat)
            r_t = self.reset_gate_bn(r_t)
            r_t = torch.sigmoid(r_t).reshape(1, B, C, H, W)
            
            # 准备候选隐状态的输入
            reset_hidden = r_t * h_state
            candidate_input = torch.cat([x_t, reset_hidden], dim=2)
            candidate_input_flat = candidate_input.flatten(0, 1)
            
            # 计算候选隐状态
            h_candidate = self.candidate_conv(candidate_input_flat)
            h_candidate = self.candidate_bn(h_candidate)
            h_candidate = torch.tanh(h_candidate).reshape(1, B, C, H, W)
            
            # 更新隐状态
            h_state = (1 - z_t) * h_state + z_t * h_candidate
            
            # 将时间差分特征与GRU状态融合
            combined = h_state + delta_t
            spike_out = self.lif(combined)
            outputs.append(spike_out)
        
        # 堆叠所有时间步的输出
        enhanced_features = torch.cat(outputs, dim=0)  # [T,B,C,H,W]
        
        # 最终的投影处理
        enhanced_flat = enhanced_features.flatten(0, 1)  # [T*B,C,H,W]
        enhanced_proj = self.gru_proj(enhanced_flat)
        enhanced_proj = self.gru_norm(enhanced_proj).reshape(T, B, C, H, W)
        
        # 将处理后的特征与原始特征融合
        final_features = x + enhanced_proj
        
        return final_features


class MS_Attention_RepConv_qkv_id(nn.Module):#Q，K，V的生成+自注意力机制+Repconv4_bn
    def __init__(
        self,
        dim,
        num_heads=8,
        qkv_bias=False,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
        sr_ratio=1,
    ):
        super().__init__()
        assert (
            dim % num_heads == 0
        ), f"dim {dim} should be divided by num_heads {num_heads}."
        self.dim = dim
        self.num_heads = num_heads
        self.scale = 0.125

        # 使用GRU_Attention_RepConv_qkv_id替代直接使用DeltaEncoder
        self.gru_attention = GRU_Attention(
            dim=dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            sr_ratio=sr_ratio,
        )
        
        self.head_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        self.head_lif1 = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        self.q_conv = nn.Sequential(RepConv(dim, dim, bias=False), nn.BatchNorm2d(dim))

        self.k_conv = nn.Sequential(RepConv(dim, dim, bias=False), nn.BatchNorm2d(dim))

        self.v_conv = nn.Sequential(RepConv(dim, dim, bias=False), nn.BatchNorm2d(dim))

        self.q_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")

        self.k_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")

        self.v_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")

        self.attn_lif = MultiStepLIFNode(
            tau=2.0, v_threshold=0.5, detach_reset=True, backend="cupy"
        )

        self.proj_conv = nn.Sequential(
            RepConv(dim, dim, bias=False), nn.BatchNorm2d(dim)
        )

    def forward(self, x):
        T, B, C, H, W = x.shape
        N = H * W

        # 使用轻量级GRU处理时间差分特征
        # enhanced_features = self.gru_attention(x)
        
        # 应用head_lif到增强特征上
        x = self.head_lif(x)  # 输入是T, B, C, H, W，输出是T, B, C, H, W
        enhanced_features = self.gru_attention(x)
        x = self.head_lif1(x)
        q = self.q_conv(enhanced_features.flatten(0, 1)).reshape(T, B, C, H, W)
        k = self.k_conv(enhanced_features.flatten(0, 1)).reshape(T, B, C, H, W)
        v = self.v_conv(enhanced_features.flatten(0, 1)).reshape(T, B, C, H, W)

        q = self.q_lif(q).flatten(3)
        q = (
            q.transpose(-1, -2)
            .reshape(T, B, N, self.num_heads, C // self.num_heads)
            .permute(0, 1, 3, 2, 4)
            .contiguous()
        )

        k = self.k_lif(k).flatten(3)
        k = (
            k.transpose(-1, -2)
            .reshape(T, B, N, self.num_heads, C // self.num_heads)
            .permute(0, 1, 3, 2, 4)
            .contiguous()
        )

        v = self.v_lif(v).flatten(3)
        v = (
            v.transpose(-1, -2)
            .reshape(T, B, N, self.num_heads, C // self.num_heads)
            .permute(0, 1, 3, 2, 4)
            .contiguous()
        )

        x = k.transpose(-2, -1) @ v  # @就是矩阵乘法，也就是k和v的相似度
        x = (q @ x) * self.scale

        # 标准自注意力机制
        # attn = (q @ k.transpose(-2, -1)) * self.scale
        # x = attn @ v
        
        x = x.transpose(3, 4).reshape(T, B, C, N).contiguous()
        x = self.attn_lif(x).reshape(T, B, C, H, W)
        x = x.reshape(T, B, C, H, W)
        x = x.flatten(0, 1)
        x = self.proj_conv(x).reshape(T, B, C, H, W)

        return x


class MS_Block(nn.Module):# Transformer-based SNN Block
    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        sr_ratio=1,
    ):
        super().__init__()

        self.attn = MS_Attention_RepConv_qkv_id(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=drop,
            sr_ratio=sr_ratio,
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()# 不做任何操作
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = MS_MLP(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop)

    def forward(self, x):
        x = x + self.attn(x)
        x = x + self.mlp(x)

        return x


class MS_DownSampling(nn.Module):
    def __init__(
        self,
        in_channels=2,
        embed_dims=256,
        kernel_size=3,
        stride=2,
        padding=1,
        first_layer=True,
    ):
        super().__init__()

        self.encode_conv = nn.Conv2d(
            in_channels,
            embed_dims,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
        )

        self.encode_bn = nn.BatchNorm2d(embed_dims)
        if not first_layer:
            self.encode_lif = MultiStepLIFNode(
                tau=2.0, detach_reset=True, backend="cupy"
            )
    
    def forward(self, x):
        T, B, _, _, _ = x.shape

        if hasattr(self, "encode_lif"):
            x = self.encode_lif(x)# (T, B, C, H, W)   
        x = self.encode_conv(x.flatten(0, 1))# (T, B, C, H, W)  (T * B, C, H, W)  (T * B, embed_dims, H', W')新高度和宽度
        _, _, H, W = x.shape# (T * B, embed_dims, H', W')   (T * B, embed_dims, H, W)
        x = self.encode_bn(x).reshape(T, B, -1, H, W).contiguous()#(T * B, embed_dims, H, W)(T, B, embed_dims, H, W)

        return x


class Spiking_vit_MetaFormer(nn.Module):
    def __init__(
        self,
        img_size_h=128,
        img_size_w=128,
        patch_size=16,
        in_channels=2,
        num_classes=11,
        embed_dim=[64, 128, 256],
        num_heads=[1, 2, 4],
        mlp_ratios=[4, 4, 4],
        qkv_bias=False,
        qk_scale=None,
        drop_rate=0.0,
        attn_drop_rate=0.0,
        drop_path_rate=0.0,
        norm_layer=nn.LayerNorm,
        depths=[6, 8, 6],
        sr_ratios=[8, 4, 2],
        kd=False,
    ):
        super().__init__()
        self.num_classes = num_classes
        self.depths = depths
        self.T = 1
        # embed_dim = [64, 128, 256, 512]

        prd = [
            x.item() for x in torch.linspace(0, drop_path_rate, depths)
        ]  # stochastic depth decay rule

        self.downsample1_1 = MS_DownSampling(
            in_channels=in_channels,
            embed_dims=embed_dim[0] // 2,
            kernel_size=7,
            stride=2,
            padding=3,
            first_layer=True,
        )

        self.ConvBlock1_1 = nn.ModuleList(
            [MS_ConvBlock(dim=embed_dim[0] // 2, mlp_ratio=mlp_ratios)]
        )

        self.downsample1_2 = MS_DownSampling(
            in_channels=embed_dim[0] // 2,
            embed_dims=embed_dim[0],
            kernel_size=3,
            stride=2,
            padding=1,
            first_layer=False,
        )

        self.ConvBlock1_2 = nn.ModuleList(
            [MS_ConvBlock(dim=embed_dim[0], mlp_ratio=mlp_ratios)]
        )

        self.downsample2 = MS_DownSampling(
            in_channels=embed_dim[0],
            embed_dims=embed_dim[1],
            kernel_size=3,
            stride=2,
            padding=1,
            first_layer=False,
        )

        self.ConvBlock2_1 = nn.ModuleList(
            [MS_ConvBlock(dim=embed_dim[1], mlp_ratio=mlp_ratios)]
        )

        self.ConvBlock2_2 = nn.ModuleList(
            [MS_ConvBlock(dim=embed_dim[1], mlp_ratio=mlp_ratios)]
        )

        self.downsample3 = MS_DownSampling(
            in_channels=embed_dim[1],
            embed_dims=embed_dim[2],
            kernel_size=3,
            stride=2,
            padding=1,
            first_layer=False,
        )

        self.block3 = nn.ModuleList(
            [
                MS_Block(
                    dim=embed_dim[2],
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratios,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop_rate,
                    attn_drop=attn_drop_rate,
                    drop_path=prd[j],
                    norm_layer=norm_layer,
                    sr_ratio=sr_ratios,
                )
                # for j in range(6)
                for j in range(4)
            ]
        )

        self.downsample4 = MS_DownSampling(
            in_channels=embed_dim[2],
            embed_dims=embed_dim[3],
            kernel_size=3,
            stride=1,
            padding=1,
            first_layer=False,
        )

        self.block4 = nn.ModuleList(
            [
                MS_Block(
                    dim=embed_dim[3],
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratios,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop_rate,
                    attn_drop=attn_drop_rate,
                    drop_path=prd[j],
                    norm_layer=norm_layer,
                    sr_ratio=sr_ratios,
                )
                for j in range(2)
            ]
        )

        self.lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend="cupy")
        self.head = (
            nn.Linear(embed_dim[3], num_classes) if num_classes > 0 else nn.Identity()
        )

        self.kd = kd
        if self.kd:
            self.head_kd = (
                nn.Linear(embed_dim[3], num_classes)
                if num_classes > 0
                else nn.Identity()
            )
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward_features(self, x):
        x = self.downsample1_1(x)
        for blk in self.ConvBlock1_1:
            x = blk(x)
        x = self.downsample1_2(x)
        for blk in self.ConvBlock1_2:
            x = blk(x)

        x = self.downsample2(x)
        for blk in self.ConvBlock2_1:
            x = blk(x)
        for blk in self.ConvBlock2_2:
            x = blk(x)

        x = self.downsample3(x)
        for blk in self.block3:
            x = blk(x)

        x = self.downsample4(x)
        for blk in self.block4:
            x = blk(x)
        return x  # T,B,C,N

    def forward(self, x):
        x = (x.unsqueeze(0)).repeat(self.T, 1, 1, 1, 1)
        x = self.forward_features(x)
        x = x.flatten(3).mean(3)
        x_lif = self.lif(x)
        x = self.head(x_lif).mean(0)
        if self.kd:
            x_kd = self.head_kd(x_lif).mean(0)
            if self.training:
                return x, x_kd
            else:
                return (x + x_kd) / 2
        return x


def metaspikformer_8_384(**kwargs):
    model = Spiking_vit_MetaFormer(
        img_size_h=224,
        img_size_w=224,
        patch_size=16,
        embed_dim=[96, 192, 384, 480],
        num_heads=8,
        mlp_ratios=4,
        in_channels=3,
        num_classes=1000,
        qkv_bias=False,
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        depths=8,
        sr_ratios=1,
        **kwargs,
    )
    return model


def metaspikformer_8_512(**kwargs):
    model = Spiking_vit_MetaFormer(
        img_size_h=224,
        img_size_w=224,
        patch_size=16,
        embed_dim=[128, 256, 512, 640],
        num_heads=8,
        mlp_ratios=4,
        in_channels=3,
        num_classes=1000,
        qkv_bias=False,
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        depths=8,
        sr_ratios=1,
        **kwargs,
    )
    return model


def metaspikformer_8_768(**kwargs):
    model = Spiking_vit_MetaFormer(
        img_size_h=224,
        img_size_w=224,
        patch_size=16,
        embed_dim=[192, 384, 768, 960],
        num_heads=8,
        mlp_ratios=4,
        in_channels=3,
        num_classes=1000,
        qkv_bias=False,
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        depths=8,
        sr_ratios=1,
        **kwargs,
    )
    return model


from timm.models import create_model

